//
//  Prim.swift
//  Chapter3LinkedList
//
//  Created by TRS-Mobile on 2021/9/6.
//

import Foundation

class Prim<T: Hashable> {
    typealias Graph = AdjacencyList<T>
    init() {
    }
    
    func produceMinimumSpanningTree(for graph: Graph) -> (cost: Double, mst: Graph) {
        var cost = 0.0
        let mst = Graph()
        var visited: Set<Vertex<T>> = []
        var priorityQueue = PriorityQueue<Edge<T>> {
            $0.weight ?? 0 < $1.weight ?? 0
        }
        
        copyVertices(from: graph, to: mst)
        guard let start = graph.vertices.first else { return (cost, mst) }
        visited.insert(start)
        addAvailableEdge(for: start, in: graph, check: visited, to: &priorityQueue)
        while let smallestEdge = priorityQueue.dequeue() {
            let vertex = smallestEdge.destination
            guard !visited.contains(vertex) else {
                continue
            }
            visited.insert(vertex)
            cost += smallestEdge.weight ?? 0
            mst.add(.undirected, from: smallestEdge.source, to: smallestEdge.destination, weight: smallestEdge.weight)
            addAvailableEdge(for: vertex, in: graph, check: visited, to: &priorityQueue)
        }
        return (cost, mst)
    }
    
    fileprivate func copyVertices(from origin: Graph, to graph: Graph) {
        for vertex in origin.vertices {
            graph.createVertex(data: vertex.data)
        }
    }
    
    fileprivate func addAvailableEdge(for vertex: Vertex<T>,
                                      in graph: Graph,
                                      check visited: Set<Vertex<T>>,
                                      to priorityQueue: inout PriorityQueue<Edge<T>>) {
        for edge in graph.edges(from: vertex) {
            if !visited.contains(edge.destination) {
                priorityQueue.enqueue(edge)
            }
        }
    }
}

